cmake_minimum_required(VERSION 3.12)
project(deeplab_tensorrt)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# Find ROS packages
find_package(catkin REQUIRED COMPONENTS
  roscpp
  rospy
  std_msgs
  sensor_msgs
  cv_bridge
  image_transport
)

# Find other dependencies
find_package(CUDA REQUIRED)
find_package(OpenCV REQUIRED)

# Allow user to pass TENSORRT_ROOT, otherwise default to a common install path
set(TENSORRT_ROOT $ENV{TENSORRT_ROOT} CACHE PATH "Path to TensorRT root")
if(NOT TENSORRT_ROOT)
    set(TENSORRT_ROOT "$ENV{HOME}/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6/TensorRT-8.5.1.7")
endif()

# Catkin package
catkin_package(
  INCLUDE_DIRS ${TENSORRT_ROOT}/include
  LIBRARIES deeplab_trt_lib
  CATKIN_DEPENDS roscpp rospy std_msgs sensor_msgs cv_bridge image_transport
  DEPENDS CUDA OpenCV
)

# Include directories
include_directories(
  ${catkin_INCLUDE_DIRS}
  ${OpenCV_INCLUDE_DIRS}
  ${TENSORRT_ROOT}/include
  ${CUDA_INCLUDE_DIRS}
)

# Link directories
link_directories(
  ${TENSORRT_ROOT}/lib
  ${CUDA_TOOLKIT_ROOT_DIR}/lib64
)

# Build standalone executable (original)
add_executable(trt_infer src/main.cpp)
target_link_libraries(trt_infer
  ${OpenCV_LIBS}
  nvinfer
  nvinfer_plugin
  ${CUDA_CUDART_LIBRARY}
)

# Build ROS node
add_executable(deeplab_trt_node src/ros_node.cpp)
target_link_libraries(deeplab_trt_node
  ${catkin_LIBRARIES}
  ${OpenCV_LIBS}
  nvinfer
  nvinfer_plugin
  ${CUDA_CUDART_LIBRARY}
)

# Set RPATH for both executables
set_target_properties(trt_infer PROPERTIES
  BUILD_RPATH "${TENSORRT_ROOT}/lib;${CUDA_TOOLKIT_ROOT_DIR}/lib64"
  INSTALL_RPATH "${TENSORRT_ROOT}/lib;${CUDA_TOOLKIT_ROOT_DIR}/lib64"
)

set_target_properties(deeplab_trt_node PROPERTIES
  BUILD_RPATH "${TENSORRT_ROOT}/lib;${CUDA_TOOLKIT_ROOT_DIR}/lib64"
  INSTALL_RPATH "${TENSORRT_ROOT}/lib;${CUDA_TOOLKIT_ROOT_DIR}/lib64"
)

# Install executables
install(TARGETS trt_infer deeplab_trt_node
  RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
)

# Install launch files
install(DIRECTORY launch/
  DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION}/launch
  FILES_MATCHING PATTERN "*.launch"
)
